import argparse

import pandas as pd
from setfit import SetFitModel

parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", type=str)
args = parser.parse_args()

df = pd.read_excel(args.input)
model = SetFitModel.from_pretrained(
    "./setfit-bge-base-v1.5-sst2-8-shot/"
)  # Load from a local directory

texts = df["问题"].tolist()

# Performing inference
preds = model.predict(texts)

df["setfit问题类别"] = preds
df.to_excel(args.input, index=False)
